{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6d9c74f7",
   "metadata": {},
   "source": [
    "# "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31d6ef5a",
   "metadata": {},
   "source": [
    "# 推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c18ebd86",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal, Optional\n",
    "\n",
    "import numpy as np\n",
    "import onnx\n",
    "import onnxruntime\n",
    "# import pytest\n",
    "from onnx import ModelProto, TensorProto, helper\n",
    "\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relax\n",
    "from tvm.relax.frontend.onnx import from_onnx\n",
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm.script import ir as I\n",
    "\n",
    "bg = np.random.MT19937(0)\n",
    "rg = np.random.Generator(bg)\n",
    "\n",
    "def generate_random_inputs(\n",
    "    model: ModelProto, inputs: Optional[dict[str, np.ndarray]] = None\n",
    ") -> dict[str, np.ndarray]:\n",
    "    input_values = {}\n",
    "    # Iterate through model inputs and extract their shape.\n",
    "    for i in model.graph.input:\n",
    "        if inputs is not None and i.name in inputs and inputs[i.name] is not None:\n",
    "            input_values[i.name] = inputs[i.name]\n",
    "            continue\n",
    "        shape = []\n",
    "        for dim in i.type.tensor_type.shape.dim:\n",
    "            shape.append(dim.dim_value)\n",
    "\n",
    "        input_values[i.name] = generate_random_value(shape, i.type.tensor_type.elem_type)\n",
    "\n",
    "    return input_values\n",
    "\n",
    "\n",
    "def generate_random_value(shape, elem_type) -> np.ndarray:\n",
    "    # Extract datatype for the input.\n",
    "    if elem_type:\n",
    "        dtype = str(helper.tensor_dtype_to_np_dtype(elem_type))\n",
    "    else:\n",
    "        dtype = \"float32\"\n",
    "\n",
    "    # Generate random inputs for each input.\n",
    "    if dtype == \"bool\":\n",
    "        # random_value = np.random.choice(a=[False, True], size=shape)\n",
    "        random_value = rg.choice(a=[False, True], size=shape)\n",
    "    elif dtype.startswith(\"int\"):\n",
    "        # Keep non-zero values\n",
    "        random_value = rg.integers(low=-63, high=63, size=shape).astype(dtype)\n",
    "        random_value[random_value <= 0] -= 1\n",
    "    else:\n",
    "        random_value = rg.standard_normal(size=shape).astype(dtype)\n",
    "\n",
    "    return random_value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8c9389c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_correctness(\n",
    "    model: ModelProto,\n",
    "    inputs: Optional[dict[str, np.ndarray]] = None,\n",
    "    ir_version: int = 8,\n",
    "    opset: int = 14,\n",
    "    rtol: float = 1e-7,\n",
    "    atol: float = 1e-5,\n",
    "    check_dtypes: bool = False,\n",
    ") -> None:\n",
    "    \"\"\"Run an onnx model in both onnxruntime and TVM through our importer\n",
    "       confirm that the results match. Otherwise, an exception will be raised.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model: ModelProto\n",
    "        The input onnx model that should be tested.\n",
    "    inputs: Optional[Dict[str, np.ndarray]]\n",
    "        An optional dictionary containing values for each input in the onnx model.\n",
    "    ir_version: int\n",
    "        Which version of the onnx IR to use.\n",
    "    opset: int\n",
    "        The opset version to use for the onnx importer.\n",
    "    atol: float\n",
    "        Set the tolerance of correctness checking. Some ops may be show more\n",
    "        arithmetic variance than others.\n",
    "    check_dtypes: bool\n",
    "        Check if data types are the same.\n",
    "    \"\"\"\n",
    "    # Configure model format.\n",
    "    if ir_version is not None:\n",
    "        model.ir_version = ir_version\n",
    "    if opset is not None:\n",
    "        model.opset_import[0].version = opset\n",
    "\n",
    "    # If inputs are not provided, extract them from the onnx graph and produce random\n",
    "    # values that we'll use for testing.\n",
    "    inputs = generate_random_inputs(model, inputs)\n",
    "\n",
    "    # Run the model through onnx to get the expected result.\n",
    "    ort_session = onnxruntime.InferenceSession(\n",
    "        model.SerializeToString(), providers=[\"CPUExecutionProvider\"]\n",
    "    )\n",
    "    ort_output = ort_session.run([], inputs)\n",
    "\n",
    "    # Convert the onnx model into relax through the onnx importer.\n",
    "    tvm_model = from_onnx(model, opset=opset, keep_params_in_input=True)\n",
    "    # Convert operators for inference mode.\n",
    "    tvm_model = relax.transform.DecomposeOpsForInference()(tvm_model)\n",
    "    # Legalize any relax ops into tensorir.\n",
    "    tvm_model = relax.transform.LegalizeOps()(tvm_model)\n",
    "\n",
    "    # Separate model from parameters.\n",
    "    tvm_model, params = relax.frontend.detach_params(tvm_model)\n",
    "    # Compile the relax graph into a VM then run.\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        ex = tvm.compile(tvm_model, target=\"llvm\")\n",
    "        vm = tvm.runtime.vm.VirtualMachine(ex, tvm.cpu())\n",
    "    # Prepare inputs.\n",
    "    input_list = [\n",
    "        inputs[key.name_hint] for key in tvm_model[\"main\"].params if key.name_hint in inputs\n",
    "    ]\n",
    "    if params:\n",
    "        input_list += params[\"main\"]\n",
    "\n",
    "    # Run model and check outputs.\n",
    "    vm.set_input(\"main\", *input_list)\n",
    "    vm.invoke_stateful(\"main\")\n",
    "    tvm_output = vm.get_outputs(\"main\")\n",
    "    # Wrap as a list if there is only one output.\n",
    "    if len(ort_output) == 1:\n",
    "        # Do not check the output number for TVM\n",
    "        # As for sequence output, the TVM output is a Tuple\n",
    "        # while the ONNX output number is one, which is a list\n",
    "        tvm_output = [tvm_output]\n",
    "\n",
    "    def _get_numpy_subdtype(narray):\n",
    "        if np.issubdtype(narray.dtype, np.integer):\n",
    "            return \"integer\"\n",
    "        elif np.issubdtype(narray.dtype, np.floating):\n",
    "            return \"floating\"\n",
    "        elif np.issubdtype(narray.dtype, np.bool_):\n",
    "            return \"bool\"\n",
    "        elif np.issubdtype(narray.dtype, np.complexfloating):\n",
    "            return \"complexfloating\"\n",
    "        else:\n",
    "            return \"other\"\n",
    "\n",
    "    def _check_output(tvm_out, ort_out):\n",
    "        if isinstance(tvm_out, tuple) and isinstance(ort_out, (tvm.runtime.ShapeTuple, list)):\n",
    "            assert len(tvm_out) == len(ort_out), \"Unequal number of outputs\"\n",
    "            for tvm_out_i, ort_out_i in zip(tvm_out, ort_out):\n",
    "                _check_output(tvm_out_i, ort_out_i)\n",
    "        elif isinstance(tvm_out, tvm.nd.NDArray) and isinstance(ort_out, np.ndarray):\n",
    "            if check_dtypes:\n",
    "                assert tvm_out.numpy().dtype == ort_out.dtype\n",
    "            tvm.testing.assert_allclose(tvm_out.numpy(), ort_out, rtol=rtol, atol=atol)\n",
    "        elif isinstance(tvm_out, tvm.runtime.ShapeTuple) and isinstance(ort_out, np.ndarray):\n",
    "            shape_out = tvm.nd.array([int(i) for i in tvm_out])\n",
    "            if check_dtypes:\n",
    "                assert _get_numpy_subdtype(shape_out.numpy()) == _get_numpy_subdtype(ort_out)\n",
    "            tvm.testing.assert_allclose(shape_out.numpy(), ort_out, rtol=rtol, atol=atol)\n",
    "        elif isinstance(tvm_out, (int, float, bool)) and isinstance(ort_out, np.ndarray):\n",
    "            if check_dtypes:\n",
    "                assert _get_numpy_subdtype(np.array(tvm_out)) == _get_numpy_subdtype(ort_out)\n",
    "            tvm.testing.assert_allclose(np.array(tvm_out), ort_out, rtol=rtol, atol=atol)\n",
    "        else:\n",
    "            raise ValueError(f\"Unsupported types: {type(tvm_out)}, {type(ort_out)}\")\n",
    "\n",
    "    # Check that number of outputs match.\n",
    "    assert len(tvm_output) == len(ort_output), \"Unequal number of outputs\"\n",
    "    for tvm_out, ort_out in zip(tvm_output, ort_output):\n",
    "        # TODO Allow configurable tolerance.\n",
    "        if ort_out is not None:\n",
    "            _check_output(tvm_out, ort_out)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "75c99fbb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "二维输入输出形状: torch.Size([32, 32])\n",
      "四维输入输出形状: torch.Size([1, 3, 4, 4])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# ===== 情况 1：二维输入，最后一维归一化 =====\n",
    "x = torch.randn(32, 32)  # 模拟 ONNX 中的 [32, 32] 输入\n",
    "scale = torch.randn(32)  # γ\n",
    "bias = torch.randn(32)   # β\n",
    "\n",
    "# 创建 LayerNorm，normalized_shape 对应最后一维大小\n",
    "layer_norm = nn.LayerNorm(normalized_shape=32, eps=1e-12, elementwise_affine=False)\n",
    "\n",
    "# 手动应用 γ 和 β（因为 elementwise_affine=False）\n",
    "y = layer_norm(x) * scale + bias\n",
    "print(\"二维输入输出形状:\", y.shape)\n",
    "\n",
    "# ===== 情况 2：四维输入，从 axis=1 开始归一化 =====\n",
    "x_img = torch.randn(1, 3, 4, 4)  # 模拟 [1, 3, 4, 4]\n",
    "\n",
    "class LayerNorm(torch.nn.Module):\n",
    "    def __init__(self, shape):\n",
    "        super().__init__()\n",
    "        self.shape = shape\n",
    "        self.weight = torch.nn.Parameter(torch.randn(shape))  # γ\n",
    "        self.bias = torch.nn.Parameter(torch.randn(shape)) # β\n",
    "\n",
    "    def forward(self, x):\n",
    "        return torch.nn.functional.layer_norm(x, self.shape, self.weight, self.bias, 1e-5)\n",
    "\n",
    "layer_norm_nd = LayerNorm((3, 4, 4))\n",
    "y_img = layer_norm_nd(x_img)\n",
    "print(\"四维输入输出形状:\", y_img.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5187b3fd",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[torch.onnx] Obtain model graph for `LayerNorm()` with `torch.export.export(..., strict=False)`...\n",
      "[torch.onnx] Obtain model graph for `LayerNorm()` with `torch.export.export(..., strict=False)`... ✅\n",
      "[torch.onnx] Run decomposition...\n",
      "[torch.onnx] Run decomposition... ✅\n",
      "[torch.onnx] Translate the graph into ONNX...\n",
      "[torch.onnx] Translate the graph into ONNX... ✅\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ONNXProgram(\n",
       "    model=\n",
       "        <\n",
       "            ir_version=10,\n",
       "            opset_imports={'': 20},\n",
       "            producer_name='pytorch',\n",
       "            producer_version='2.9.0.dev20250725+cpu',\n",
       "            domain=None,\n",
       "            model_version=None,\n",
       "        >\n",
       "        graph(\n",
       "            name=main_graph,\n",
       "            inputs=(\n",
       "                %\"x\"<FLOAT,[1,3,4,4]>\n",
       "            ),\n",
       "            outputs=(\n",
       "                %\"layer_norm\"<FLOAT,[1,3,4,4]>\n",
       "            ),\n",
       "            initializers=(\n",
       "                %\"weight\"<FLOAT,[3,4,4]>{TorchTensor(...)},\n",
       "                %\"bias\"<FLOAT,[3,4,4]>{TorchTensor(...)}\n",
       "            ),\n",
       "        ) {\n",
       "            0 |  # node_layer_norm\n",
       "                 %\"layer_norm\"<FLOAT,[1,3,4,4]>, %\"\"<?,?>, %\"\"<?,?> ⬅️ ::LayerNormalization(%\"x\", %\"weight\"{...}, %\"bias\"{...}) {axis=-3, epsilon=1e-05, stash_type=1}\n",
       "            return %\"layer_norm\"<FLOAT,[1,3,4,4]>\n",
       "        }\n",
       "\n",
       "\n",
       "    ,\n",
       "    exported_program=\n",
       "        ExportedProgram:\n",
       "            class GraphModule(torch.nn.Module):\n",
       "                def forward(self, p_weight: \"f32[3, 4, 4]\", p_bias: \"f32[3, 4, 4]\", x: \"f32[1, 3, 4, 4]\"):\n",
       "                     # File: /tmp/ipykernel_1344764/1027669049.py:27 in forward, code: return torch.nn.functional.layer_norm(x, self.shape, self.weight, self.bias, 1e-5)\n",
       "                    layer_norm: \"f32[1, 3, 4, 4]\" = torch.ops.aten.layer_norm.default(x, [3, 4, 4], p_weight, p_bias);  x = p_weight = p_bias = None\n",
       "                    return (layer_norm,)\n",
       "            \n",
       "        Graph signature: \n",
       "            # inputs\n",
       "            p_weight: PARAMETER target='weight'\n",
       "            p_bias: PARAMETER target='bias'\n",
       "            x: USER_INPUT\n",
       "    \n",
       "            # outputs\n",
       "            layer_norm: USER_OUTPUT\n",
       "    \n",
       "        Range constraints: {}\n",
       "\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "temp_dir = Path(\".temp\")\n",
    "temp_dir.mkdir(exist_ok=True)\n",
    "torch.onnx.export(\n",
    "    layer_norm_nd, \n",
    "    (x_img,), \n",
    "    temp_dir/'demo.onnx', \n",
    "    input_names=[\"x\",],\n",
    "    dynamo=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "99abfd05",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm/python/tvm/relax/frontend/onnx/onnx_frontend.py:3949: UserWarning: No Op registered for LayerNormalization with domain_version of 14\n",
      "\n",
      "==> Context: Bad node spec for node. Name: node_layer_norm OpType: LayerNormalization\n",
      "  warnings.warn(str(exception))\n"
     ]
    }
   ],
   "source": [
    "check_correctness(\n",
    "    model=onnx.load(temp_dir/'demo.onnx'),\n",
    "    inputs=[{\"x\": x_img,}],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "61dc07cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
       "<span style=\"color: #007979; font-style: italic\"># from tvm.script import relax as R</span>\n",
       "\n",
       "<span style=\"color: #A2F\">@I</span><span style=\"color: #A2F; font-weight: bold\">.</span>ir_module\n",
       "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #00F; font-weight: bold\">Module</span>:\n",
       "    <span style=\"color: #A2F\">@R</span><span style=\"color: #A2F; font-weight: bold\">.</span>function\n",
       "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #00F\">main</span>(x: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), weight: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>), bias: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>)) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>):\n",
       "        R<span style=\"color: #A2F; font-weight: bold\">.</span>func_attr({<span style=\"color: #BA2121\">&quot;num_input&quot;</span>: <span style=\"color: #008000\">1</span>, <span style=\"color: #BA2121\">&quot;params&quot;</span>: [metadata[<span style=\"color: #BA2121\">&quot;ffi.NDArray&quot;</span>][<span style=\"color: #008000\">0</span>], metadata[<span style=\"color: #BA2121\">&quot;ffi.NDArray&quot;</span>][<span style=\"color: #008000\">1</span>]]})\n",
       "        <span style=\"color: #008000; font-weight: bold\">with</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>dataflow():\n",
       "            lv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> R<span style=\"color: #A2F; font-weight: bold\">.</span>nn<span style=\"color: #A2F; font-weight: bold\">.</span>layer_norm(x, weight, bias, axes<span style=\"color: #A2F; font-weight: bold\">=</span>[<span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">3</span>, <span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">2</span>, <span style=\"color: #A2F; font-weight: bold\">-</span><span style=\"color: #008000\">1</span>], epsilon<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">9.9999997473787516e-06</span>, center<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>, scale<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>)\n",
       "            gv: R<span style=\"color: #A2F; font-weight: bold\">.</span>Tensor((<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">3</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">4</span>), dtype<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;float32&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">=</span> lv\n",
       "            R<span style=\"color: #A2F; font-weight: bold\">.</span>output(gv)\n",
       "        <span style=\"color: #008000; font-weight: bold\">return</span> gv\n",
       "\n",
       "<span style=\"color: #007979; font-style: italic\"># Metadata omitted. Use show_meta=True in script() method to show it.</span>\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "tvm_model = from_onnx(onnx.load(temp_dir/'demo.onnx'), opset=20, keep_params_in_input=True)\n",
    "tvm_model.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "904d5613",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
